#!/usr/bin/env python3
import argparse
import os
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Set

from elftools.elf.elffile import ELFFile

try:
    from capstone import Cs, CS_ARCH_ARM64, CS_MODE_LITTLE_ENDIAN
    from capstone.arm64 import (
        ARM64_INS_ADR,
        ARM64_INS_ADRP,
        ARM64_INS_ADD,
        ARM64_INS_LDR,
        ARM64_INS_STR,
        ARM64_INS_MOV,
        ARM64_INS_MOVZ,
        ARM64_INS_MOVK,
        ARM64_INS_BL,
        ARM64_INS_BLR,
        ARM64_INS_RET,
        ARM64_REG_SP,
        ARM64_OP_IMM,
        ARM64_OP_MEM,
        ARM64_OP_REG,
    )
except Exception as e:
    raise SystemExit("Missing capstone. Install with: pip install capstone or apt install python3-capstone")


def sign_extend(val: int, bits: int) -> int:
    sign_bit = 1 << (bits - 1)
    return (val & (sign_bit - 1)) - (val & sign_bit)


def decode_adrp_page_base(insn32: int, pc: int) -> Optional[int]:
    try:
        immlo = (insn32 >> 29) & 0x3
        immhi = (insn32 >> 5) & 0x7FFFF
        imm = (immhi << 2) | immlo
        imm = sign_extend(imm, 21) << 12
        page = pc & ~0xFFF
        return (page + imm) & 0xFFFFFFFFFFFFFFFF
    except Exception:
        return None


def decode_adr_target(insn32: int, pc: int) -> Optional[int]:
    # ADR immediate uses same immhi/immlo fields but not shifted by 12, and PC not page-aligned
    try:
        immlo = (insn32 >> 29) & 0x3
        immhi = (insn32 >> 5) & 0x7FFFF
        imm = (immhi << 2) | immlo
        imm = sign_extend(imm, 21)
        return (pc + imm) & 0xFFFFFFFFFFFFFFFF
    except Exception:
        return None


def load_symbols(elf: ELFFile) -> Dict[str, int]:
    syms: Dict[str, int] = {}
    for sec_name in (".symtab", ".dynsym"):
        sec = elf.get_section_by_name(sec_name)
        if not sec:
            continue
        try:
            for sym in sec.iter_symbols():
                name = sym.name
                val = int(sym.entry.get("st_value", 0) or 0)
                if name and val:
                    syms[name] = val
        except Exception:
            continue
    return syms


def collect_exec_sections(elf: ELFFile) -> List[Tuple[int, bytes]]:
    res: List[Tuple[int, bytes]] = []
    for sec in elf.iter_sections():
        hdr = sec.header
        sh_flags = int(hdr.get("sh_flags", 0))
        sh_addr = int(hdr.get("sh_addr", 0))
        sh_type = hdr.get("sh_type")
        if (sh_flags & 0x4) != 0:  # SHF_EXECINSTR
            try:
                data = sec.data()
            except Exception:
                continue
            if data:
                res.append((sh_addr, data))
    # Sort by address to create a linear view
    res.sort(key=lambda x: x[0])
    return res


@dataclass
class TargetSpec:
    name: str
    max_params: int


@dataclass
class ParamHit:
    index: int
    reg: str
    value: Optional[int]
    string: Optional[str]


@dataclass
class CallHit:
    target: str
    callsite: int
    params: List[ParamHit]


MAX_REGISTER_PARAMS = 8


class Insn:
    __slots__ = ("addr", "id", "ops", "bytes")

    def __init__(self, addr: int, iid: int, ops, b: bytes):
        self.addr = addr
        self.id = iid
        self.ops = ops
        self.bytes = b


def disasm_aarch64(elf: ELFFile) -> List[Insn]:
    md = Cs(CS_ARCH_ARM64, CS_MODE_LITTLE_ENDIAN)
    md.detail = True
    insns: List[Insn] = []
    for base, data in collect_exec_sections(elf):
        for i in md.disasm(data, base):
            insns.append(Insn(i.address, i.id, i.operands, bytes(i.bytes)))
    insns.sort(key=lambda x: x.addr)
    return insns


def build_plt_map(elf: ELFFile) -> Dict[str, int]:
    plt = elf.get_section_by_name('.plt')
    rela_plt = elf.get_section_by_name('.rela.plt')
    dynsym = elf.get_section_by_name('.dynsym')
    if not plt or not rela_plt or not dynsym:
        return {}
    base = int(plt.header.get('sh_addr', 0))
    # AArch64: first 0x20 bytes are resolver stub, each slot 0x10 bytes
    slot0 = base + 0x20
    stride = 0x10
    name_to_addr: Dict[str, int] = {}
    idx = 0
    for rel in rela_plt.iter_relocations():
        try:
            sym_idx = rel.entry['r_info_sym']
            sym = dynsym.get_symbol(sym_idx)
            name = sym.name
        except Exception:
            name = ''
        if name:
            name_to_addr[name] = slot0 + idx * stride
        idx += 1
    return name_to_addr


def normalize_symbol_name(name: str) -> str:
    base = name
    if base.endswith('@plt'):
        base = base[:-4]
    if '@@' in base:
        base = base.split('@@', 1)[0]
    if '@' in base:
        base = base.split('@', 1)[0]
    if base.endswith('.plt'):
        base = base[:-4]
    return base


def collect_target_addresses(elf: ELFFile, symtab: Dict[str, int], target_names: List[str]) -> Dict[str, Set[int]]:
    targets: Dict[str, Set[int]] = {name: set() for name in target_names}
    def maybe_add(symbol: str, va: int) -> None:
        if not va:
            return
        base = normalize_symbol_name(symbol)
        for t in target_names:
            if base == t:
                targets[t].add(va)
    for name, va in symtab.items():
        maybe_add(name, va)
    pltmap = build_plt_map(elf)
    for name, va in pltmap.items():
        maybe_add(name, va)
    return targets


def va_to_file_off(elf: ELFFile, va: int) -> Optional[int]:
    for seg in elf.iter_segments():
        if seg["p_type"] != "PT_LOAD":
            continue
        vaddr = int(seg["p_vaddr"]) ; memsz = int(seg["p_memsz"]) ; filesz = int(seg["p_filesz"]) ; off = int(seg["p_offset"]) 
        if va >= vaddr and va < vaddr + memsz:
            rel = va - vaddr
            if rel < filesz:
                return off + rel
            return None
    return None


def read_c_string_at_va(elf: ELFFile, va: int, max_len: int = 4096) -> Optional[str]:
    off = va_to_file_off(elf, va)
    if off is None:
        return None
    s = elf.stream
    s.seek(off)
    data = bytearray()
    for _ in range(max_len):
        b = s.read(1)
        if not b:
            break
        if b == b"\x00":
            break
        data.extend(b)
    try:
        return data.decode("utf-8")
    except Exception:
        try:
            return data.decode("latin-1")
        except Exception:
            return None


def resolve_regs_before(insns: List[Insn], idx: int, target_regs: List[str]) -> Dict[str, Optional[int]]:
    reg_val: Dict[str, int] = {}
    reg_parts: Dict[str, Dict[int, int]] = {}
    reg_page_addr: Dict[str, int] = {}
    stack_slots: Dict[int, int] = {}
    pending_sp_loads: List[Tuple[str, int]] = []
    pending_sp_store: Dict[int, str] = {}
    pending_add_imm: Dict[str, int] = {}
    pending_add_dep: Dict[str, Tuple[str, int]] = {}
    targets = set(target_regs)

    md = Cs(CS_ARCH_ARM64, CS_MODE_LITTLE_ENDIAN)

    def reg_name(rid: int) -> str:
        if rid == ARM64_REG_SP:
            return 'sp'
        name = md.reg_name(rid)  # type: ignore
        if name.startswith('w') and name[1:].isdigit():
            return 'x' + name[1:]
        return name

    def fulfill_with_reg(rname: str) -> None:
        if rname in reg_val:
            for d, (b, imm2) in list(pending_add_dep.items()):
                if b == rname:
                    reg_val[d] = reg_val[b] + imm2
                    del pending_add_dep[d]
        for off, sreg in list(pending_sp_store.items()):
            if sreg == rname and rname in reg_val:
                stack_slots[off] = reg_val[rname]
                for (d2, o2) in list(pending_sp_loads):
                    if o2 == off:
                        reg_val[d2] = stack_slots[off]
                        pending_sp_loads.remove((d2, o2))
                del pending_sp_store[off]

    def done() -> bool:
        return all(reg in reg_val for reg in targets) and not pending_sp_loads and not pending_add_dep

    start = max(0, idx - 80)
    for i in range(idx - 1, start - 1, -1):
        if done():
            break
        ins = insns[i]
        iid = ins.id
        ops = ins.ops
        pc = ins.addr

        if iid in (ARM64_INS_RET, ARM64_INS_BL, ARM64_INS_BLR):
            break

        if iid == ARM64_INS_ADRP and len(ops) >= 2 and ops[0].type == ARM64_OP_REG:
            dst = reg_name(ops[0].reg)
            insn32 = int.from_bytes(ins.bytes[:4], 'little')
            base = decode_adrp_page_base(insn32, pc)
            if base is not None:
                reg_page_addr[dst] = base
                if dst in pending_add_imm:
                    reg_val[dst] = base + pending_add_imm[dst]
                    fulfill_with_reg(dst)
                fulfill_with_reg(dst)
            continue

        if iid == ARM64_INS_ADR and len(ops) >= 2 and ops[0].type == ARM64_OP_REG:
            dst = reg_name(ops[0].reg)
            insn32 = int.from_bytes(ins.bytes[:4], 'little')
            addr = decode_adr_target(insn32, pc)
            if addr is not None:
                reg_val[dst] = addr
                fulfill_with_reg(dst)
            continue

        if iid == ARM64_INS_ADD and len(ops) >= 3 and ops[0].type == ARM64_OP_REG and ops[1].type == ARM64_OP_REG and ops[2].type == ARM64_OP_IMM:
            dst = reg_name(ops[0].reg)
            base = reg_name(ops[1].reg)
            imm = int(ops[2].imm)
            if base in reg_val:
                reg_val[dst] = reg_val[base] + imm
                fulfill_with_reg(dst)
            elif base in reg_page_addr:
                reg_val[dst] = reg_page_addr[base] + imm
                fulfill_with_reg(dst)
            else:
                if dst == base:
                    pending_add_imm[dst] = imm
                else:
                    pending_add_dep[dst] = (base, imm)
            continue

        if iid == ARM64_INS_LDR and len(ops) >= 2 and ops[0].type == ARM64_OP_REG and ops[1].type == ARM64_OP_MEM:
            dst = reg_name(ops[0].reg)
            base_r = ops[1].mem.base
            base = reg_name(base_r) if base_r != 0 else ''
            off = int(ops[1].mem.disp or 0)
            if base in reg_val:
                reg_val[dst] = reg_val[base] + off
                fulfill_with_reg(dst)
            elif base == 'sp':
                if off in stack_slots:
                    reg_val[dst] = stack_slots[off]
                else:
                    pending_sp_loads.append((dst, off))
            continue

        if iid == ARM64_INS_STR and len(ops) >= 2 and ops[0].type == ARM64_OP_REG and ops[1].type == ARM64_OP_MEM:
            src = reg_name(ops[0].reg)
            base_r = ops[1].mem.base
            base = reg_name(base_r) if base_r != 0 else ''
            off = int(ops[1].mem.disp or 0)
            if base == 'sp':
                if src in reg_val:
                    stack_slots[off] = reg_val[src]
                    for (d2, o2) in list(pending_sp_loads):
                        if o2 == off:
                            reg_val[d2] = stack_slots[off]
                            pending_sp_loads.remove((d2, o2))
                else:
                    pending_sp_store[off] = src
            continue

        if iid in (ARM64_INS_MOVZ, ARM64_INS_MOVK) and len(ops) >= 2 and ops[0].type == ARM64_OP_REG and ops[1].type == ARM64_OP_IMM:
            dst = reg_name(ops[0].reg)
            imm = int(ops[1].imm)
            try:
                sh = getattr(ops[1].shift, 'value', 0) or 0
            except Exception:
                sh = 0
            part = (imm & 0xFFFF) << sh
            if iid == ARM64_INS_MOVZ:
                reg_parts[dst] = {0: 0, 16: 0, 32: 0, 48: 0}
                reg_parts[dst][sh] = (part >> sh) & 0xFFFF
                val = 0
                for k, v in reg_parts[dst].items():
                    val |= (v & 0xFFFF) << k
                reg_val[dst] = val
            else:
                if dst not in reg_parts:
                    reg_parts[dst] = {0: 0, 16: 0, 32: 0, 48: 0}
                reg_parts[dst][sh] = (part >> sh) & 0xFFFF
                val = 0
                for k, v in reg_parts[dst].items():
                    val |= (v & 0xFFFF) << k
                reg_val[dst] = val
            fulfill_with_reg(dst)
            continue

        if iid == ARM64_INS_MOV and len(ops) >= 2 and ops[0].type == ARM64_OP_REG:
            dst = reg_name(ops[0].reg)
            if ops[1].type == ARM64_OP_REG:
                src = reg_name(ops[1].reg)
                if src in reg_val:
                    reg_val[dst] = reg_val[src]
                    fulfill_with_reg(dst)
            elif ops[1].type == ARM64_OP_IMM:
                reg_val[dst] = int(ops[1].imm)
                fulfill_with_reg(dst)
            continue

    return {reg: reg_val.get(reg) for reg in target_regs}


def analyze(path: str, specs: List[TargetSpec]) -> List[CallHit]:
    with open(path, 'rb') as f:
        elf = ELFFile(f)
        mach = elf.header.get('e_machine')
        if not (mach == 'EM_AARCH64' or mach == 183):
            raise SystemExit(f"ELF machine '{mach}' is not AArch64")

        if not specs:
            return []

        symbols = load_symbols(elf)
        target_names = [spec.name for spec in specs]
        addr_sets = collect_target_addresses(elf, symbols, target_names)
        addr_to_name: Dict[int, str] = {}
        for name, addrs in addr_sets.items():
            for va in addrs:
                addr_to_name.setdefault(va, name)

        if not addr_to_name:
            return []

        spec_by_name = {spec.name: spec for spec in specs}

        insns = disasm_aarch64(elf)

        results: List[CallHit] = []
        for i, ins in enumerate(insns):
            if ins.id != ARM64_INS_BL or len(ins.ops) < 1 or ins.ops[0].type != ARM64_OP_IMM:
                continue
            tgt = int(ins.ops[0].imm)
            target_name = addr_to_name.get(tgt)
            if not target_name:
                continue
            spec = spec_by_name.get(target_name)
            if not spec:
                continue
            param_count = min(spec.max_params, MAX_REGISTER_PARAMS)
            regs = [f'x{n}' for n in range(param_count)]
            resolved = resolve_regs_before(insns, i, regs) if regs else {}
            params: List[ParamHit] = []
            for idx, reg in enumerate(regs):
                val = resolved.get(reg) if resolved else None
                sval = read_c_string_at_va(elf, val) if isinstance(val, int) else None
                params.append(ParamHit(index=idx + 1, reg=reg, value=val, string=sval))
            results.append(CallHit(target=target_name, callsite=ins.addr, params=params))
        return results


def parse_function_specs(values: List[str], default_limit: int) -> List[TargetSpec]:
    if default_limit < 0:
        raise SystemExit('--param-limit must be non-negative')
    specs: List[TargetSpec] = []
    for item in values:
        if not item:
            continue
        if ':' in item:
            name, limit_str = item.split(':', 1)
            if not name:
                raise SystemExit(f"Invalid --function spec: {item}")
            try:
                limit = int(limit_str) if limit_str else default_limit
            except ValueError:
                raise SystemExit(f"Invalid parameter limit in --function {item}")
        else:
            name = item
            limit = default_limit
        if limit < 0:
            raise SystemExit(f"Parameter limit must be >=0 for {name}")
        specs.append(TargetSpec(name=name, max_params=limit))
    if not specs:
        specs.append(TargetSpec(name='dlsym', max_params=default_limit))
    return specs


def main():
    ap = argparse.ArgumentParser(description='Inspect selected function calls in AArch64 ELF files using Capstone')
    ap.add_argument('elf', help='Path to AArch64 ELF shared object (.so)')
    ap.add_argument('-f', '--function', action='append', dest='functions', default=[], help='Function to inspect, optionally with :N to cap parameters (e.g. dlsym:2). Repeatable.')
    ap.add_argument('--param-limit', type=int, default=2, help='Default number of parameters to resolve per function (default: 2)')
    args = ap.parse_args()

    if not os.path.isfile(args.elf):
        raise SystemExit(f"Not a file: {args.elf}")

    specs = parse_function_specs(args.functions, args.param_limit)

    try:
        results = analyze(args.elf, specs)
    except Exception as e:
        raise SystemExit(f"Error: {e}")

    print(f"File: {args.elf}")
    if not results:
        print('No matching callsites found.')
        return
    for hit in results:
        if hit.params:
            formatted = []
            for param in hit.params:
                if param.value is None:
                    val_s = '?'
                else:
                    val_s = f"0x{param.value:x}"
                if param.string is not None:
                    quoted = param.string.replace('"', '\"')
                    val_s += f' -> "{quoted}"'
                formatted.append(f"#{param.index} {param.reg}={val_s}")
            params_str = ', '.join(formatted)
        else:
            params_str = '(no params tracked)'
        print(f"- 0x{hit.callsite:x}: {hit.target}({params_str})")


if __name__ == '__main__':
    main()
